"""
Template Builder Utility

Processes expert videos to create action templates with:
- Normalized skeleton sequences
- DINO feature vectors
- Geometric metrics (ROM, speed, etc.)
- Phase annotations
"""

import os
import yaml
import numpy as np
from typing import Dict, Optional, List
import logging

from src.core.pose_detector import PoseDetector
from src.core.skeleton_normalizer import SkeletonNormalizer, calculate_joint_angles
from src.core.feature_extractor import FeatureExtractor
from src.core.action_segmenter import PhaseDetector

logger = logging.getLogger(__name__)


class TemplateBuilder:
    """
    Builds action templates from expert demonstration videos.
    """

    def __init__(self,
                 templates_dir: str = "./templates",
                 config: Optional[Dict] = None,
                 pose_detector: Optional[PoseDetector] = None,
                 skeleton_normalizer: Optional[SkeletonNormalizer] = None,
                 feature_extractor: Optional[FeatureExtractor] = None):
        """
        Initialize template builder.

        Args:
            templates_dir: Base directory for templates
            config: System configuration dictionary
            pose_detector: Optional PoseDetector instance (reuses existing if provided)
            skeleton_normalizer: Optional SkeletonNormalizer instance (reuses existing if provided)
            feature_extractor: Optional FeatureExtractor instance (reuses existing if provided)
        """
        self.templates_dir = templates_dir
        self.config = config or {}

        # Use provided instances or create new ones
        if pose_detector is not None:
            self.pose_detector = pose_detector
        else:
            pose_config = self.config.get('pose_detection', {})
            self.pose_detector = PoseDetector(
                model_complexity=pose_config.get('model_complexity', 1),
                min_detection_confidence=pose_config.get('min_detection_confidence', 0.5),
                min_tracking_confidence=pose_config.get('min_tracking_confidence', 0.5)
            )

        if skeleton_normalizer is not None:
            self.skeleton_normalizer = skeleton_normalizer
        else:
            skeleton_config = self.config.get('skeleton', {})
            self.skeleton_normalizer = SkeletonNormalizer(
                canvas_size=tuple(skeleton_config.get('canvas_size', [224, 224])),
                scale_fill=skeleton_config.get('scale_fill', 0.85),
                line_thickness=skeleton_config.get('line_thickness', 2)
            )

        if feature_extractor is not None:
            self.feature_extractor = feature_extractor
        else:
            feature_config = self.config.get('feature_extraction', {})
            self.feature_extractor = FeatureExtractor(
                model_name=feature_config.get('model_name', 'facebook/dinov2-base'),
                device=feature_config.get('device', 'mps'),
                batch_size=feature_config.get('batch_size', 16)
            )

        self.phase_detector = PhaseDetector()

        logger.info("TemplateBuilder initialized")

    def build_template(self,
                      video_path: str,
                      action_type: str,
                      template_name: str = "expert",
                      metadata: Optional[Dict] = None) -> Dict:
        """
        Build a complete template from expert video.

        Args:
            video_path: Path to expert demonstration video
            action_type: Type of action (e.g., 'raise_leg_left')
            template_name: Name for this template (default: 'expert')
            metadata: Additional metadata (description, key joints, etc.)

        Returns:
            Template dictionary containing all processed data
        """
        logger.info(f"Building template: {action_type}/{template_name} from {video_path}")

        # Step 1: Pose detection
        logger.info("Step 1: Detecting poses...")
        landmarks_sequence, video_metadata = self.pose_detector.detect_from_video(video_path)

        # Interpolate missing frames if any
        missing_count = sum(1 for x in landmarks_sequence if x is None)
        if missing_count > 0:
            logger.warning(f"Found {missing_count} missing frames, interpolating...")
            landmarks_sequence = self.pose_detector.interpolate_missing_frames(landmarks_sequence)

        # Step 2: Normalize and render skeletons
        logger.info("Step 2: Normalizing and rendering skeletons...")
        normalized_sequence, rendered_images = self.skeleton_normalizer.process_sequence(
            landmarks_sequence, render=True
        )

        # Step 3: Extract DINO features
        logger.info("Step 3: Extracting DINO features...")
        features = self.feature_extractor.extract_from_sequence(rendered_images)

        # Step 4: Calculate geometric metrics
        logger.info("Step 4: Calculating geometric metrics...")
        geometric_metrics = self._calculate_geometric_metrics(normalized_sequence, video_metadata['fps'])

        # Step 5: Detect phases
        logger.info("Step 5: Detecting action phases...")
        phases = self.phase_detector.detect_phases(normalized_sequence, video_metadata['fps'])

        # Create template dictionary
        template = {
            'action_type': action_type,
            'template_name': template_name,
            'video_path': video_path,
            'num_frames': len(normalized_sequence),
            'fps': video_metadata['fps'],
            'features_shape': features.shape,
            'geometric_metrics': geometric_metrics,
            'phases': phases,
            'metadata': metadata or {},
            'created_at': video_metadata.get('created_at', 'unknown')
        }

        # Save template
        template_dir = os.path.join(self.templates_dir, action_type)
        os.makedirs(template_dir, exist_ok=True)

        # Save features
        features_path = os.path.join(template_dir, f"{template_name}_features.npy")
        np.save(features_path, features)
        template['features_path'] = features_path

        # Save normalized landmarks
        landmarks_path = os.path.join(template_dir, f"{template_name}_landmarks.npy")
        landmarks_array = np.array(normalized_sequence)
        np.save(landmarks_path, landmarks_array)
        template['landmarks_path'] = landmarks_path

        # Save metadata
        meta_path = os.path.join(template_dir, f"{template_name}_meta.yaml")
        with open(meta_path, 'w', encoding='utf-8') as f:
            # Convert numpy types to Python native types for YAML
            yaml_safe_template = self._convert_to_yaml_safe(template)
            yaml.dump(yaml_safe_template, f, allow_unicode=True, default_flow_style=False)
        template['meta_path'] = meta_path

        logger.info(f"Template built successfully: {meta_path}")
        logger.info(f"  - Features: {features_path}")
        logger.info(f"  - Landmarks: {landmarks_path}")

        return template

    def _calculate_geometric_metrics(self,
                                    normalized_sequence: list,
                                    fps: float) -> Dict:
        """
        Calculate geometric metrics for template.

        Args:
            normalized_sequence: Normalized landmark sequence
            fps: Frame rate

        Returns:
            Dictionary of geometric metrics
        """
        metrics = {}

        # Calculate angles for all frames
        angles_sequence = []
        for landmarks in normalized_sequence:
            angles = calculate_joint_angles(landmarks)
            angles_sequence.append(angles)

        # ROM (Range of Motion) for each joint
        rom = {}
        for joint in angles_sequence[0].keys():
            angles = [frame_angles[joint] for frame_angles in angles_sequence]
            rom[joint] = {
                'min': float(np.min(angles)),
                'max': float(np.max(angles)),
                'range': float(np.max(angles) - np.min(angles)),
                'mean': float(np.mean(angles)),
                'std': float(np.std(angles))
            }
        metrics['rom'] = rom

        # Calculate velocities
        positions = np.array([lm.flatten() for lm in normalized_sequence])
        velocities = np.diff(positions, axis=0) * fps
        velocity_magnitudes = np.linalg.norm(velocities, axis=1)

        metrics['speed'] = {
            'mean': float(np.mean(velocity_magnitudes)),
            'max': float(np.max(velocity_magnitudes)),
            'std': float(np.std(velocity_magnitudes))
        }

        # Calculate accelerations (for stability)
        accelerations = np.diff(velocities, axis=0) * fps
        acceleration_magnitudes = np.linalg.norm(accelerations, axis=1)

        metrics['stability'] = {
            'jitter': float(np.sqrt(np.mean(accelerations ** 2))),
            'smoothness': float(1.0 / (1.0 + np.std(acceleration_magnitudes)))
        }

        # Left-right symmetry
        left_right_pairs = [
            ('left_elbow', 'right_elbow'),
            ('left_knee', 'right_knee'),
            ('left_hip', 'right_hip')
        ]

        symmetry = {}
        for left_joint, right_joint in left_right_pairs:
            left_angles = [frame[left_joint] for frame in angles_sequence]
            right_angles = [frame[right_joint] for frame in angles_sequence]
            diff = np.abs(np.array(left_angles) - np.array(right_angles))
            symmetry[f"{left_joint}_{right_joint}"] = {
                'mean_diff': float(np.mean(diff)),
                'max_diff': float(np.max(diff)),
                'symmetry_score': float(100 * np.exp(-np.mean(diff) / 30.0))
            }

        metrics['symmetry'] = symmetry

        return metrics

    def _convert_to_yaml_safe(self, data):
        """Convert numpy types to Python native types for YAML serialization."""
        if isinstance(data, dict):
            return {key: self._convert_to_yaml_safe(value) for key, value in data.items()}
        elif isinstance(data, (list, tuple)):
            # Convert both list and tuple to list for YAML
            return [self._convert_to_yaml_safe(item) for item in data]
        elif isinstance(data, np.integer):
            return int(data)
        elif isinstance(data, np.floating):
            return float(data)
        elif isinstance(data, np.ndarray):
            return data.tolist()
        else:
            return data

    def load_template(self, action_type: str, template_name: str = "expert") -> Dict:
        """
        Load a pre-built template.

        Args:
            action_type: Type of action
            template_name: Template name

        Returns:
            Template dictionary
        """
        template_dir = os.path.join(self.templates_dir, action_type)

        # Load metadata
        meta_path = os.path.join(template_dir, f"{template_name}_meta.yaml")
        with open(meta_path, 'r', encoding='utf-8') as f:
            template = yaml.safe_load(f)

        # Load features
        features_path = template['features_path']
        template['features'] = np.load(features_path)

        # Load landmarks
        landmarks_path = template['landmarks_path']
        template['landmarks'] = np.load(landmarks_path)

        logger.info(f"Loaded template: {action_type}/{template_name}")

        return template

    def list_templates(self) -> Dict[str, list]:
        """
        List all available templates.

        Returns:
            Dictionary mapping action_type to list of template names
        """
        templates = {}

        if not os.path.exists(self.templates_dir):
            return templates

        for action_type in os.listdir(self.templates_dir):
            action_dir = os.path.join(self.templates_dir, action_type)
            if not os.path.isdir(action_dir):
                continue

            template_names = []
            for file in os.listdir(action_dir):
                if file.endswith('_meta.yaml'):
                    template_name = file.replace('_meta.yaml', '')
                    template_names.append(template_name)

            if template_names:
                templates[action_type] = template_names

        return templates

    def load_all_templates(self, template_name: str = 'expert') -> Dict[str, Dict]:
        """
        Load all templates for all action types.

        Args:
            template_name: Name of template version to load (default: 'expert')

        Returns:
            Dictionary mapping action_type to template data
        """
        all_templates = {}
        template_list = self.list_templates()

        for action_type, names in template_list.items():
            if template_name in names:
                try:
                    template = self.load_template(action_type, template_name)
                    all_templates[action_type] = template
                    logger.info(f"Loaded template: {action_type}/{template_name}")
                except Exception as e:
                    logger.warning(f"Failed to load template {action_type}/{template_name}: {e}")

        logger.info(f"Loaded {len(all_templates)} templates: {list(all_templates.keys())}")
        return all_templates

    def build_templates_from_segments(self,
                                     video_path: str,
                                     segments_labels: Dict[int, str],
                                     template_name: str = 'expert',
                                     metadata: Optional[Dict] = None) -> Dict[str, Dict]:
        """
        Build multiple templates from a long video with labeled segments.

        This allows creating multiple action templates from a single expert video
        that contains several different actions.

        Args:
            video_path: Path to expert video file
            segments_labels: Dict mapping segment_id to action_type
                           e.g., {0: 'squat', 1: 'leg_raise', 2: 'squat'}
            template_name: Name for these templates (default: 'expert')
            metadata: Additional metadata

        Returns:
            Dictionary of {action_type: template_paths} for created templates

        Example:
            >>> builder = TemplateBuilder()
            >>> # Video has 3 segments detected
            >>> labels = {0: 'squat', 1: 'leg_raise', 2: 'squat'}
            >>> templates = builder.build_templates_from_segments(
            ...     'expert_multi.mp4', labels)
        """
        logger.info(f"Building templates from segments in {video_path}")
        logger.info(f"Segment labels: {segments_labels}")

        # Step 1: Process full video
        logger.info("Step 1: Detecting poses...")
        landmarks_seq, video_metadata = self.pose_detector.detect_from_video(video_path)
        landmarks_seq = self.pose_detector.interpolate_missing_frames(landmarks_seq)

        logger.info("Step 2: Normalizing and rendering skeletons...")
        normalized_seq, rendered_imgs = self.skeleton_normalizer.process_sequence(landmarks_seq)

        logger.info("Step 3: Extracting DINO features...")
        features = self.feature_extractor.extract_from_sequence(rendered_imgs)

        # Step 4: Segment the video
        logger.info("Step 4: Segmenting video...")
        from ..core.action_segmenter import ActionSegmenter
        segmenter = ActionSegmenter(
            fps=video_metadata.get('fps', 30.0),
            k_high=1.0,
            k_low=0.5,
            L_on=3,
            L_off=6,
            min_duration=0.5,  # Minimum 0.5s to filter noise
            min_silence=0.3,  # Minimum 0.3s between actions
            adaptive=False  # Disable adaptive mode
        )
        segments = segmenter.segment_sequence(normalized_seq, action_type='default')

        if len(segments) != len(segments_labels):
            logger.warning(f"Detected {len(segments)} segments but got {len(segments_labels)} labels")
            logger.warning(f"Detected segments: {segments}")

        # Step 5: Build template for each labeled segment
        built_templates = {}

        for seg_id, action_type in segments_labels.items():
            if seg_id >= len(segments):
                logger.warning(f"Segment {seg_id} not found (only {len(segments)} segments detected)")
                continue

            start, end = segments[seg_id]
            logger.info(f"Building template for segment {seg_id} (frames {start}-{end}): {action_type}")

            # Extract segment data
            seg_landmarks = normalized_seq[start:end]
            seg_features = features[start:end]
            seg_rendered = rendered_imgs[start:end]

            # Build template for this segment
            # Create temp metadata similar to full video
            temp_metadata = {
                'fps': video_metadata.get('fps', 30.0),
                'num_frames': end - start,
                'video_path': video_path,
                'segment_range': [int(start), int(end)]
            }
            if metadata:
                temp_metadata.update(metadata)

            # Use internal method to build and save template
            template_paths = self._save_template_data(
                action_type=action_type,
                template_name=template_name,
                landmarks_seq=seg_landmarks,
                features=seg_features,
                metadata=temp_metadata
            )

            built_templates[f"{action_type}_{seg_id}"] = template_paths
            logger.info(f"✓ Built template: {action_type}/{template_name} (segment {seg_id})")

        logger.info(f"Built {len(built_templates)} templates from {len(segments_labels)} labeled segments")
        return built_templates

    def _save_template_data(self,
                           action_type: str,
                           template_name: str,
                           landmarks_seq: List[np.ndarray],
                           features: np.ndarray,
                           metadata: Dict) -> Dict:
        """
        Internal method to save template data.

        Args:
            action_type: Type of action
            template_name: Name of template
            landmarks_seq: Normalized landmarks
            features: DINO features
            metadata: Template metadata

        Returns:
            Dictionary with paths to saved files
        """
        # Create action directory
        action_dir = os.path.join(self.templates_dir, action_type)
        os.makedirs(action_dir, exist_ok=True)

        # File paths
        meta_path = os.path.join(action_dir, f"{template_name}_meta.yaml")
        features_path = os.path.join(action_dir, f"{template_name}_features.npy")
        landmarks_path = os.path.join(action_dir, f"{template_name}_landmarks.npy")

        # Calculate geometric metrics
        from ..core.skeleton_normalizer import calculate_joint_angles
        fps = metadata.get('fps', 30.0)
        geometric_metrics = self._calculate_geometric_metrics(landmarks_seq, fps)

        # Detect phases
        from ..core.action_segmenter import PhaseDetector
        phase_detector = PhaseDetector()
        phases = phase_detector.detect_phases(landmarks_seq, metadata.get('fps', 30.0))

        # Save features
        np.save(features_path, features)

        # Save landmarks
        landmarks_array = np.array(landmarks_seq)
        np.save(landmarks_path, landmarks_array)

        # Save metadata
        template_meta = {
            'action_type': action_type,
            'template_name': template_name,
            'num_frames': len(landmarks_seq),
            'features_shape': list(features.shape),
            'features_path': features_path,
            'landmarks_path': landmarks_path,
            'geometric_metrics': geometric_metrics,
            'phases': phases,
            'created_at': metadata.get('created_at', 'unknown'),
            'fps': metadata.get('fps', 30.0),
            'video_path': metadata.get('video_path', 'unknown'),
            'metadata': metadata
        }

        with open(meta_path, 'w') as f:
            yaml.dump(self._convert_to_yaml_safe(template_meta), f)

        return {
            'meta_path': meta_path,
            'features_path': features_path,
            'landmarks_path': landmarks_path
        }
